#include "compi/channel.h"

#include <gtest/gtest.h>

#include <mpi.h>
#include <thread>
#include <vector>

using namespace compi;

class ChannelTest : public ::testing::Test
{
protected:
    void SetUp() override
    {
        // 共享全局 Environment，避免在每个用例结束时析构导致 MPI_Finalize
        int initialized = 0;
        MPI_Initialized(&initialized);
        if (!initialized)
        {
            int argc = 0;
            char ** argv = nullptr;
            MPI_Init(&argc, &argv);
        }
        static Environment * shared_env = nullptr;
        if (!shared_env)
            shared_env = new Environment();
        env_ = shared_env; // 借用，不拥有
    }

    void TearDown() override
    {
        // 清理所有通道
        ChannelManager::for_comm(MPI_COMM_WORLD).cleanup_all();
    }

protected:
    Environment * env_{};
};

// 测试基本的通道创建和销毁
TEST_F(ChannelTest, BasicChannelCreation)
{
    ChannelId id("test_channel");
    auto & manager = ChannelManager::for_comm(MPI_COMM_WORLD);

    // 创建通道
    auto create_res = manager.create_channel(id, true);
    ASSERT_FALSE(create_res.is_error());
    auto channel = create_res.value();
    ASSERT_NE(channel, nullptr);
    EXPECT_EQ(channel->id().name, "test_channel");
    EXPECT_EQ(channel->state(), ChannelState::Created);

    // 激活通道
    EXPECT_TRUE(channel->activate().is_success());
    EXPECT_EQ(channel->state(), ChannelState::Active);

    // 验证通道属性
    EXPECT_NE(channel->comm(), MPI_COMM_NULL);
    EXPECT_GE(channel->rank(), 0);
    EXPECT_GT(channel->size(), 0);

    // 关闭通道
    EXPECT_TRUE(channel->close().is_success());
    EXPECT_EQ(channel->state(), ChannelState::Closed);
}

// 测试通道管理器的基本功能
TEST_F(ChannelTest, ChannelManagerBasics)
{
    auto & manager = ChannelManager::for_comm(MPI_COMM_WORLD);
    ChannelId id1("channel1");
    ChannelId id2("channel2");

    // 创建两个通道
    auto res1 = manager.create_channel(id1, true);
    auto res2 = manager.create_channel(id2, true);
    ASSERT_TRUE(res1.is_success());
    ASSERT_TRUE(res2.is_success());
    auto channel1 = res1.value();
    auto channel2 = res2.value();
    ASSERT_NE(channel1, nullptr);
    ASSERT_NE(channel2, nullptr);
    EXPECT_NE(channel1, channel2);

    // 通过ID获取通道
    auto retrieved1 = manager.get_channel(id1);
    auto retrieved2 = manager.get_channel(id2);

    EXPECT_EQ(channel1, retrieved1);
    EXPECT_EQ(channel2, retrieved2);

    // 激活通道并检查计数
    EXPECT_TRUE(channel1->activate().is_success());
    EXPECT_TRUE(channel2->activate().is_success());
    EXPECT_EQ(manager.active_channel_count(), 2);

    // 销毁一个通道
    EXPECT_TRUE(manager.destroy_channel(id1).is_success());
    EXPECT_EQ(manager.get_channel(id1), nullptr);
    EXPECT_NE(manager.get_channel(id2), nullptr);
}

// 测试通道的独立性
TEST_F(ChannelTest, ChannelIsolation)
{
    auto & manager = ChannelManager::for_comm(MPI_COMM_WORLD);
    ChannelId id1("isolated_channel1");
    ChannelId id2("isolated_channel2");

    auto res1 = manager.create_channel(id1, true);
    auto res2 = manager.create_channel(id2, true);
    ASSERT_TRUE(res1.is_success());
    ASSERT_TRUE(res2.is_success());
    auto channel1 = res1.value();
    auto channel2 = res2.value();

    ASSERT_NE(channel1, nullptr);
    ASSERT_NE(channel2, nullptr);

    EXPECT_TRUE(channel1->activate().is_success());
    EXPECT_TRUE(channel2->activate().is_success());

    // 验证通道有不同的通信器
    EXPECT_NE(channel1->comm(), channel2->comm());
    EXPECT_NE(channel1->comm(), MPI_COMM_WORLD);
    EXPECT_NE(channel2->comm(), MPI_COMM_WORLD);

    // 验证通道属性相同（因为都基于MPI_COMM_WORLD）
    EXPECT_EQ(channel1->rank(), channel2->rank());
    EXPECT_EQ(channel1->size(), channel2->size());
}

// 测试通道组功能
TEST_F(ChannelTest, ChannelGroupBasics)
{
    auto & manager = ChannelManager::for_comm(MPI_COMM_WORLD);

    // 创建通道组（直接构造）
    ChannelGroup group("test_group");
    EXPECT_EQ(group.size(), 0);

    // 创建通道并添加到组
    ChannelId id1("group_channel1");
    ChannelId id2("group_channel2");

    auto res1 = manager.create_channel(id1, true);
    auto res2 = manager.create_channel(id2, true);
    ASSERT_TRUE(res1.is_success());
    ASSERT_TRUE(res2.is_success());
    auto channel1 = res1.value();
    auto channel2 = res2.value();

    EXPECT_TRUE(group.add_channel(channel1));
    EXPECT_TRUE(group.add_channel(channel2));
    EXPECT_EQ(group.size(), 2);

    // 从组中获取通道
    auto retrieved = group.get_channel(id1);
    EXPECT_EQ(retrieved, channel1);

    // 激活组中所有通道
    EXPECT_EQ(group.activate_all(), 2);
    EXPECT_EQ(channel1->state(), ChannelState::Active);
    EXPECT_EQ(channel2->state(), ChannelState::Active);

    // 关闭组中所有通道
    EXPECT_EQ(group.close_all(), 2);
    EXPECT_EQ(channel1->state(), ChannelState::Closed);
    EXPECT_EQ(channel2->state(), ChannelState::Closed);

    // 从组中移除通道
    EXPECT_TRUE(group.remove_channel(id1));
    EXPECT_EQ(group.size(), 1);
    EXPECT_EQ(group.get_channel(id1), nullptr);
}

// 测试并发安全性
TEST_F(ChannelTest, ConcurrentAccess)
{
    auto & manager = ChannelManager::for_comm(MPI_COMM_WORLD);
    const int num_threads = 4;
    const int channels_per_thread = 10;
    std::vector<std::thread> threads;
    std::vector<std::vector<std::shared_ptr<Channel>>> thread_channels(num_threads);

    // 并发创建通道
    for (int t = 0; t < num_threads; ++t)
    {
        threads.emplace_back(
            [&, t]()
            {
                for (int i = 0; i < channels_per_thread; ++i)
                {
                    ChannelId id("concurrent_channel_" + std::to_string(t) + "_" +
                                 std::to_string(i));
                    auto res = manager.create_channel(id, true);
                    if (res.is_success())
                    {
                        auto channel = res.value();
                        if (channel->activate().is_success())
                            thread_channels[t].push_back(channel);
                    }
                }
            });
    }

    // 等待所有线程完成
    for (auto & thread : threads)
    {
        thread.join();
    }

    // 验证所有通道都创建成功
    int total_channels = 0;
    for (const auto & channels : thread_channels)
    {
        total_channels += channels.size();
    }
    EXPECT_EQ(total_channels, num_threads * channels_per_thread);
    EXPECT_EQ(manager.active_channel_count(), total_channels);
}

// 测试资源管理和清理
TEST_F(ChannelTest, ResourceManagement)
{
    auto & manager = ChannelManager::for_comm(MPI_COMM_WORLD);

    // 创建多个通道
    std::vector<std::shared_ptr<Channel>> channels;
    for (int i = 0; i < 5; ++i)
    {
        ChannelId id("resource_channel_" + std::to_string(i));
        auto res = manager.create_channel(id, true);
        ASSERT_TRUE(res.is_success());
        auto channel = res.value();
        ASSERT_NE(channel, nullptr);
        ASSERT_TRUE(channel->activate().is_success());
        channels.push_back(channel);
    }

    EXPECT_EQ(manager.active_channel_count(), 5);

    // 手动关闭一些通道
    channels[0]->close();
    channels[2]->close();

    // 清理所有通道
    manager.cleanup_all();

    // 验证所有通道都被清理
    EXPECT_EQ(manager.active_channel_count(), 0);

    // 验证通过ID无法再获取到通道
    for (int i = 0; i < 5; ++i)
    {
        ChannelId id("resource_channel_" + std::to_string(i));
        EXPECT_EQ(manager.get_channel(id), nullptr);
    }
}

// 测试通道状态转换
TEST_F(ChannelTest, StateTransitions)
{
    ChannelId id("state_test_channel");
    auto & manager = ChannelManager::for_comm(MPI_COMM_WORLD);
    auto create_res2 = manager.create_channel(id, true);
    ASSERT_TRUE(create_res2.is_success());
    auto channel = create_res2.value();

    ASSERT_NE(channel, nullptr);
    EXPECT_EQ(channel->state(), ChannelState::Created);

    // Created -> Active
    EXPECT_TRUE(channel->activate().is_success());
    EXPECT_EQ(channel->state(), ChannelState::Active);

    // 重复激活应该失败
    EXPECT_TRUE(channel->activate().is_error());
    EXPECT_EQ(channel->state(), ChannelState::Active);

    // Active -> Closed
    EXPECT_TRUE(channel->close().is_success());
    EXPECT_EQ(channel->state(), ChannelState::Closed);

    // 重复关闭应该成功
    EXPECT_TRUE(channel->close().is_success());
    EXPECT_EQ(channel->state(), ChannelState::Closed);

    // 关闭后无法激活
    EXPECT_TRUE(channel->activate().is_error());
    EXPECT_EQ(channel->state(), ChannelState::Closed);
}